#include "hnurm_armor/Compensator.h"
#include <angles/angles.h>
#include <opencv2/core.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/opencv.hpp>
#include <utility>

namespace hnurm
{
    Compensator::Compensator(rclcpp::Node::SharedPtr node)
        : node_(std::move(node))
    {
        gravity_ = static_cast<float>(node_->declare_parameter("compensator.gravity", 9.8));
        friction_coeff_ = static_cast<float>(node_->declare_parameter("compensator.friction_coeff", 0.006));
        bullet_speed_bias_ = static_cast<float>(node_->declare_parameter("compensator.bullet_speed_bias", -0.5));
        channel_delay_ = static_cast<float>(node_->declare_parameter("compensator.channel_delay", 1.2));
        yaw_bias_ = static_cast<float>(node_->declare_parameter("compensator.yaw_bias", 0.0));
        pitch_bias_ = static_cast<float>(node_->declare_parameter("compensator.pitch_bias", 0.0));
        imutofrt_y = static_cast<float>(node_->declare_parameter("compensator.imutofrt_y", 0.0));
        imutofrt_z = static_cast<float>(node_->declare_parameter("compensator.imutofrt_z", 0.0));


        RCLCPP_INFO(node_->get_logger(),
                    "Compensator initialized, yaw_bias: %f, pitch_bias: %f",
                    yaw_bias_,
                    pitch_bias_);

        fx = fy = fz = fv = yaw_difference = 0;
        pout = yout = 0;
        tmp_fly_time = channel_delay_;
        bullet_speed_ = 0;
        px = py = pz = pv = {};
    }

    float Compensator::calculateAverage(const std::vector<float> &vec)
    {
        int n = vec.size();
        if (n == 0) return 0.1f; // 如果向量为空，返回0

        float sum = 0.0f;
        for (float value: vec)
        {
            sum += value;
        }
        return sum / n;
    }
    void Compensator::Averagedistance(std::vector<float> &p, float value)
    {
        if (p.size() >= 50)
        {
            p.erase(p.begin());
        }
        p.push_back(value);
    }
    bool Compensator::cmp(const predict_armor &a, const predict_armor &b)
    {
        return a.yaw < b.yaw;
    }
    double Compensator::diff_eq(double v) const
    {
        return -gravity_ - (friction_coeff_ * abs(v) * v);
    }
    double Compensator::integrate_euler(double v0, double t)
    {
        double v = v0;
        double y = 0.0;
        double dt = 0.001;
        for (double time = 0; time <= t; time += dt)
        {
            double dv = diff_eq(v);
            v += dv * dt;
            y += v * dt;
        }
        return y;
    }
    float Compensator::CalcFinalAngle(
        TargetInfo &target_msg,
        hnurm_interfaces::msg::VisionRecvData &recv_data,
        hnurm_interfaces::msg::VisionSendData &send_data,
        ArmorsNum &armors_num
    )
    {
        tmp_fly_time = channel_delay_;
        // todo: fix this
        auto bullet_speed = recv_data.bullet_speed.data;
        auto g_yaw = recv_data.yaw;
        t_info = target_msg;
        bullet_speed_ = static_cast<float>(bullet_speed) + bullet_speed_bias_;
        bullet_speed_ = 25.5;
        if (armors_num == ArmorsNum::OUTPOST_3 || armors_num == ArmorsNum::BASE_1)
        {
            Averagedistance(px, t_info.x);
            Averagedistance(py, t_info.y);
            Averagedistance(pz, t_info.z);
            Averagedistance(pv, t_info.w_yaw);
        }
        if (armors_num != ArmorsNum::OUTPOST_3 && armors_num != ArmorsNum::BASE_1
            && sqrt(pow(t_info.vx, 2) + pow(t_info.vy, 2) + pow(t_info.vz, 2)) > 0.3)
        {
            for (int i = 0; i < 2; i++)
            {
                PredictPosition();
                CompensatePitch();
            }
        }
        for (int i = 0; i < 4; i++)
        {
            PredictPose(armors_num);
            CompensatePitch();
        }
        yout = (float) (g_yaw
            + angles::shortest_angular_distance(g_yaw * CV_PI / 180, -atan2(fx_center, fy_center)) * 180 / CV_PI);
        auto armor_yaw =
            (float) (g_yaw + angles::shortest_angular_distance(g_yaw * CV_PI / 180, -atan2(fx, fy)) * 180 / CV_PI);
        if (abs(armor_yaw - yout) > 3.3)
        {
            yaw_difference = DBL_MAX;
        }
        auto difference = yaw_difference;
        send_data.pitch = pout + pitch_bias_;
        send_data.yaw = yout + yaw_bias_;
        return difference;
    }
    void Compensator::PredictPosition()
    {
        fx = t_info.x + tmp_fly_time * t_info.vx - t_info.radius_1 * cos(t_info.yaw);
        fy = t_info.y + tmp_fly_time * t_info.vy - t_info.radius_1 * sin(t_info.yaw);
        fz = t_info.z + tmp_fly_time * t_info.vz;
    }
    void Compensator::PredictPose(ArmorsNum &armors_num)
    {
        if (abs(t_info.w_yaw) < 1)
        {
            if (armors_num == ArmorsNum::OUTPOST_3 || armors_num == ArmorsNum::BASE_1)
            {
                fx = calculateAverage(px) - t_info.radius_1 * cos(t_info.yaw);
                fy = calculateAverage(py) - t_info.radius_1 * sin(t_info.yaw);
                fz = calculateAverage(pz);
            }
            else
            {
                fx = t_info.x + tmp_fly_time * t_info.vx - t_info.radius_1 * cos(t_info.yaw);
                fy = t_info.y + tmp_fly_time * t_info.vy - t_info.radius_1 * sin(t_info.yaw);
                fz = t_info.z + tmp_fly_time * t_info.vz;
            }
            fx_center = fx;
            fy_center = fy;
            yaw_difference = 0;
        }
        else
        {
            float yaw;
            if (armors_num == ArmorsNum::OUTPOST_3 || armors_num == ArmorsNum::BASE_1)
            {
                fx = calculateAverage(px);
                fy = calculateAverage(py);
                fz = calculateAverage(pz);
                fv = calculateAverage(pv);
                yaw = t_info.yaw + tmp_fly_time * fv;
            }
            else
            {
                fx = t_info.x + tmp_fly_time * t_info.vx;
                fy = t_info.y + tmp_fly_time * t_info.vy;
                fz = t_info.z + tmp_fly_time * t_info.vz;
                yaw = t_info.yaw + tmp_fly_time * t_info.w_yaw;
            }
            fx_center = fx;
            fy_center = fy;
            float yaw_center = -atan2(fx_center, fy_center) + CV_PI / 2;
            predict_armor armor1, armor2, armor3, armor4;
            std::vector<predict_armor> armors = {};
            if (armors_num == ArmorsNum::OUTPOST_3)
            {
                armor1.position.x = fx - t_info.radius_1 * cos(yaw);
                armor1.position.y = fy - t_info.radius_1 * sin(yaw);
                armor1.position.z = fz;
                armor1.yaw = abs(angles::shortest_angular_distance(yaw, yaw_center));
                armor2.position.x = fx - t_info.radius_1 * cos(float(yaw + 2 * CV_PI / 3));
                armor2.position.y = fy - t_info.radius_1 * sin(float(yaw + 2 * CV_PI / 3));
                armor2.position.z = fz;
                armor2.yaw = abs(angles::shortest_angular_distance(yaw + 2 * CV_PI / 3, yaw_center));
                armor3.position.x = fx - t_info.radius_1 * sin(float(yaw + 4 * CV_PI / 3));
                armor3.position.y = fy - t_info.radius_1 * cos(float(yaw + 4 * CV_PI / 3));
                armor3.position.z = fz;
                armor3.yaw = abs(angles::shortest_angular_distance(yaw + 4 * CV_PI / 3, yaw_center));
                armors.push_back(armor1);
                armors.push_back(armor2);
                armors.push_back(armor3);
            }
            else if (armors_num == ArmorsNum::NORMAL_4)
            {
                armor1.position.x = fx - t_info.radius_1 * cos(yaw);
                armor1.position.y = fy - t_info.radius_1 * sin(yaw);
                armor1.position.z = fz;
                armor1.yaw = abs(angles::shortest_angular_distance(yaw, yaw_center));
                armor2.position.x = fx + t_info.radius_1 * cos(yaw);
                armor2.position.y = fy + t_info.radius_1 * sin(yaw);
                armor2.position.z = fz;
                armor2.yaw = abs(angles::shortest_angular_distance(yaw + CV_PI, yaw_center));
                armor3.position.x = fx - t_info.radius_2 * sin(yaw);
                armor3.position.y = fy + t_info.radius_2 * cos(yaw);
                armor3.position.z = fz + t_info.dz;
                armor3.yaw = abs(angles::shortest_angular_distance(yaw + 3 * CV_PI / 2, yaw_center));
                armor4.position.x = fx + t_info.radius_2 * sin(yaw);
                armor4.position.y = fy - t_info.radius_2 * cos(yaw);
                armor4.position.z = fz + t_info.dz;
                armor4.yaw = abs(angles::shortest_angular_distance(yaw + CV_PI / 2, yaw_center));
                armors.push_back(armor1);
                armors.push_back(armor2);
                armors.push_back(armor3);
                armors.push_back(armor4);
            }
            else if (armors_num == ArmorsNum::BALANCE_2)
            {
                armor1.position.x = fx - t_info.radius_1 * cos(yaw);
                armor1.position.y = fy - t_info.radius_1 * sin(yaw);
                armor1.position.z = fz;
                armor1.yaw = abs(angles::shortest_angular_distance(yaw, yaw_center));
                armor2.position.x = fx + t_info.radius_1 * cos(yaw);
                armor2.position.y = fy + t_info.radius_1 * sin(yaw);
                armor2.position.z = fz;
                armor2.yaw = abs(angles::shortest_angular_distance(yaw + CV_PI, yaw_center));
                armors.push_back(armor1);
                armors.push_back(armor2);
            }
            else if (armors_num == ArmorsNum::BASE_1)
            {
                armor1.position.x = fx - t_info.radius_1 * cos(yaw);
                armor1.position.y = fy - t_info.radius_1 * sin(yaw);
                armor1.position.z = fz;
                armors.push_back(armor1);
            }
            sort(armors.begin(), armors.end(), cmp);
            fx = armors[0].position.x;
            fy = armors[0].position.y;
            fz = armors[0].position.z;
            yaw_difference = armors[0].yaw;
            if (armors_num != ArmorsNum::BASE_1 && armors_num != ArmorsNum::OUTPOST_3 && abs(t_info.w_yaw) < 6)
            {
                fx_center = fx;
                fy_center = fy;
                yaw_difference = 0;
            }
        }
    }

    void Compensator::CompensatePitch()
    {
        double dist_horizon = sqrt(pow(fx, 2) + pow(fy - imutofrt_y, 2));
        double target_height = fz - imutofrt_z;
        // 迭代参数
        double vx, vy, fly_time, tmp_height = target_height, delta_height = target_height, tmp_pitch, real_height;
        int iteration_num = 0;
        while ((iteration_num <= 30 && abs(delta_height) >= 0.005) || iteration_num <= 5)
        {
            tmp_pitch = atan((tmp_height) / dist_horizon);
            dist_horizon = sqrt(pow(fx, 2) + pow(fy - imutofrt_y * cos(tmp_pitch) + imutofrt_z * sin(tmp_pitch), 2));
            target_height = fz - imutofrt_z * cos(tmp_pitch) - imutofrt_y * sin(tmp_pitch);
            vx = bullet_speed_ * cos(tmp_pitch);
            vy = bullet_speed_ * sin(tmp_pitch);

            fly_time = (exp(friction_coeff_ * dist_horizon) - 1) / (friction_coeff_ * vx);
            real_height = integrate_euler(vy, fly_time);
            delta_height = target_height - real_height;
            tmp_height += delta_height;
            iteration_num++;
        }
        tmp_fly_time = (float) fly_time + channel_delay_;// 计算新的飞行时间
        pout = (float) (tmp_pitch * 180 / CV_PI);        // 计算补偿角度
    }
}  // namespace hnurm